package net.neptech.config;


import com.alibaba.cloud.ai.graph.GraphRepresentation;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.OverAllStateFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.EdgeAction;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.node.QuestionClassifierNode;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import net.neptech.node.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;
import java.util.Map;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;
import static com.alibaba.cloud.ai.graph.action.AsyncEdgeAction.edge_async;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;

/**
 * 问诊工作流自动配置类
 * 定义从患者输入到最终推荐的完整问诊流程
 */
@Configuration
public class MedicalGraphConfig {

  private static final Logger logger = LoggerFactory.getLogger(MedicalGraphConfig.class);

  @Autowired
  private ChatClient chatClient;


/**
 * 创建一个名为overAllStateFactory的Bean工厂方法
 * 该方法用于创建OverAllStateFactory实例，用于生成OverAllState对象
 * @return 返回一个OverAllStateFactory函数式接口的实现，使用方法引用方式创建OverAllState实例
 */
  @Bean
  public OverAllStateFactory overAllStateFactory() {  // 定义一个名为overAllStateFactory的Bean方法
    return OverAllState::new;  // 使用方法引用返回OverAllState的构造函数引用
  }
  /**
   * 定义问诊工作流图
   * @return 配置好的工作流图
   * @throws GraphStateException 图状态异常
   */
  @Bean
  public StateGraph medicalWorkflowGraph(ApplicationContext applicationContext) throws GraphStateException {
    // 构建AI客户端，添加日志顾问
//    ChatClient chatClient = builder.defaultAdvisors(new SimpleLoggerAdvisor()).build();

    // 定义状态工厂 - 工作流的"共享内存"
    OverAllStateFactory stateFactory = () -> {
      OverAllState state = new OverAllState();
      // 注册状态键和更新策略（全部使用替换策略）
      state.registerKeyAndStrategy("input", new ReplaceStrategy());          // 用户原始输入
      state.registerKeyAndStrategy("intent", new ReplaceStrategy());         // 意图判断结果
      state.registerKeyAndStrategy("symptoms", new ReplaceStrategy());       // 提取的症状
      state.registerKeyAndStrategy("standardized_symptoms", new ReplaceStrategy()); // 标准化后的症状
      state.registerKeyAndStrategy("sensitive_flag", new ReplaceStrategy()); // 敏感词标记
      state.registerKeyAndStrategy("severity", new ReplaceStrategy());       // 严重程度
      state.registerKeyAndStrategy("disease_inference", new ReplaceStrategy()); // 疾病推断
      state.registerKeyAndStrategy("knowledge_graph_result", new ReplaceStrategy()); // 知识图谱结果
      state.registerKeyAndStrategy("need_additional_questions", new ReplaceStrategy()); // 是否需要追加问题
      state.registerKeyAndStrategy("additional_info", new ReplaceStrategy()); // 患者补充信息
      state.registerKeyAndStrategy("department_recommendation", new ReplaceStrategy()); // 科室推荐
      state.registerKeyAndStrategy("risk_level", new ReplaceStrategy());     // 风险等级
      state.registerKeyAndStrategy("final_recommendation", new ReplaceStrategy()); // 最终推荐
      return state;
    };

    // 创建各个处理节点（具体实现在node包下）
    // 1. 意图判断节点 - 区分问诊和非问诊意图
    QuestionClassifierNode intentJudgment = QuestionClassifierNode.builder()
            .chatClient(chatClient)
            .inputTextKey("input")
            .categories(List.of("medical_consultation", "general_conversation", "other_intent"))
            .classificationInstructions(List.of(
                    "请严格判断用户输入是否属于医疗问诊意图",
                    "医疗问诊意图的特征：包含症状描述（如头痛、咳嗽、发烧）、疾病名称、就医咨询、科室选择、用药咨询等",
                    "非医疗问诊意图的特征：问候语（你好、早上好）、自我介绍（你是谁）、闲聊、其他非医疗话题",
                    "如果用户描述具体症状或询问医疗建议，分类为 medical_consultation",
                    "如果只是普通对话或非医疗问题，分类为 general_conversation",
                    "示例：'我头痛发烧' -> medical_consultation, '你是谁' -> general_conversation",
                    "'咳嗽该挂什么科' -> medical_consultation, '今天天气怎么样' -> general_conversation"
            ))
            .build();

    // 2. 症状收集节点 - 从用户输入中提取症状
    QuestionClassifierNode symptomCollection = QuestionClassifierNode.builder()
            .chatClient(chatClient)
            .inputTextKey("input")
            .categories(List.of("symptoms_extracted"))
            .classificationInstructions(List.of(
                    "从用户输入中提取所有提到的症状",
                    "将每个症状单独列出"
            ))
            .build();

    // 3. 症状标准化节点 - 将症状转换为医学术语
    QuestionClassifierNode symptomStandardization = QuestionClassifierNode.builder()
            .chatClient(chatClient)
            .inputTextKey("symptoms")
            .categories(List.of("standardized_symptoms", "contains_sensitive_terms"))
            .classificationInstructions(List.of(
                    "将症状标准化为医学术语",
                    "标记是否存在需要人工审核的敏感或警报术语"
            ))
            .build();

    // 4. 严重程度分类节点 - 评估症状严重程度
    QuestionClassifierNode severityClassification = QuestionClassifierNode.builder()
            .chatClient(chatClient)
            .inputTextKey("standardized_symptoms")
            .categories(List.of("critical", "severe", "moderate", "mild"))
            .classificationInstructions(List.of(
                    "根据医疗标准对症状严重程度进行分类",
                    "将可能危及生命的情况标记为危急"
            ))
            .build();

    // 5. 疾病推理节点 - 根据症状推断可能的疾病
    QuestionClassifierNode diseaseInference = QuestionClassifierNode.builder()
            .chatClient(chatClient)
            .inputTextKey("standardized_symptoms")
            .categories(List.of("possible_diagnoses"))
            .classificationInstructions(List.of(
                    "根据症状推断可能的疾病",
                    "优先考虑最常见和最危险的可能性"
            ))
            .build();

    // 6. 知识图谱查询节点 - 查询医学知识图谱
    QuestionClassifierNode knowledgeGraphQuery = QuestionClassifierNode.builder()
            .chatClient(chatClient)
            .inputTextKey("disease_inference")
            .categories(List.of("knowledge_graph_results"))
            .classificationInstructions(List.of(
                    "查询医学知识图谱获取相关疾病和症状",
                    "返回最相关的连接和信息"
            ))
            .build();

    // 7. 鉴别诊断判断节点 - 判断是否需要追加问题
    QuestionClassifierNode differentialDiagnosisDecision = QuestionClassifierNode.builder()
            .chatClient(chatClient)
            .inputTextKey("knowledge_graph_result")
            .categories(List.of("needs_additional_questions", "sufficient_info"))
            .classificationInstructions(List.of(
                    "判断是否需要追加问题进行鉴别诊断",
                    "评估当前信息是否足以做出可靠诊断"
            ))
            .build();

    // 8. 科室推荐节点 - 推荐合适的就诊科室
    QuestionClassifierNode departmentRecommendation = QuestionClassifierNode.builder()
            .chatClient(chatClient)
            .inputTextKey("disease_inference")
            .categories(List.of("recommended_departments"))
            .classificationInstructions(List.of(
                    "根据症状和可能的诊断推荐合适的科室",
                    "考虑最可能和最严重的病情"
            ))
            .build();

    // 9. 风险复核节点 - 重新评估风险等级
    QuestionClassifierNode riskReview = QuestionClassifierNode.builder()
            .chatClient(chatClient)
            .inputTextKey("department_recommendation")
            .categories(List.of("high_risk", "low_risk"))
            .classificationInstructions(List.of(
                    "根据科室推荐重新评估风险等级",
                    "标记可能被遗漏的高风险情况"
            ))
            .build();

    // 构建工作流图
    StateGraph stateGraph = new StateGraph("在线问诊工作流", stateFactory)
            // 添加所有处理节点（使用异步执行）
            .addNode("intent_judgment", node_async(intentJudgment))          // 意图判断
            .addNode("symptom_collection", node_async(symptomCollection))   // 症状收集
            .addNode("symptom_standardization", node_async(symptomStandardization)) // 症状标准化
            .addNode("manual_review", node_async(new ManualReviewNode()))    // 人工审核
            .addNode("severity_classification", node_async(severityClassification)) // 严重程度分类
            .addNode("disease_inference", node_async(diseaseInference))      // 疾病推理
            .addNode("knowledge_graph_query", node_async(knowledgeGraphQuery)) // 知识图谱查询
            .addNode("differential_diagnosis_decision", node_async(differentialDiagnosisDecision)) // 鉴别诊断判断
            .addNode("additional_question", node_async(new AdditionalQuestionNode())) // 追加提问
            .addNode("patient_info_supplementation", node_async(new PatientInfoSupplementationNode())) // 患者补充信息
            .addNode("department_recommendation", node_async(departmentRecommendation)) // 科室推荐
            .addNode("risk_review", node_async(riskReview))                 // 风险复核
            .addNode("emergency_triage", node_async(new EmergencyTriageNode())) // 紧急分诊
            .addNode("non_consult_response", node_async(new NonConsultResponseNode(applicationContext))) // 非问诊响应
            .addNode("final_recommendation", node_async(new FinalRecommendationNode())) // 最终推荐

            // 定义流程走向
            // 1. 开始 -> 意图判断
            .addEdge(START, "intent_judgment")

            // 2. 意图判断后的分支
            .addConditionalEdges("intent_judgment",
                    edge_async(new IntentDispatcher()),
                    Map.of(
                            "consultation", "symptom_collection",  // 问诊意图 -> 症状收集
                            "non_consultation", "non_consult_response" // 非问诊意图 -> 非问诊响应
                    ))

            // 3. 症状收集 -> 症状标准化
            .addEdge("symptom_collection", "symptom_standardization")

            // 4. 症状标准化后的分支
            .addConditionalEdges("symptom_standardization",
                    edge_async(new StandardizationDispatcher()),
                    Map.of(
                            "standardized", "severity_classification", // 标准化成功 -> 严重程度分类
                            "sensitive", "manual_review"              // 含敏感词 -> 人工审核
                    ))

            // 5. 人工审核 -> 症状标准化（重新处理）
            .addEdge("manual_review", "symptom_standardization")

            // 6. 严重程度分类后的分支
            .addConditionalEdges("severity_classification",
                    edge_async(new SeverityDispatcher()),
                    Map.of(
                            "critical", "emergency_triage",  // 危急 -> 紧急分诊
                            "severe", "emergency_triage",    // 严重 -> 紧急分诊
                            "moderate", "disease_inference", // 中度 -> 疾病推理
                            "mild", "disease_inference"      // 轻度 -> 疾病推理
                    ))

            // 7. 疾病推理 -> 知识图谱查询 -> 鉴别诊断判断
            .addEdge("disease_inference", "knowledge_graph_query")
            .addEdge("knowledge_graph_query", "differential_diagnosis_decision")

            // 8. 鉴别诊断判断后的分支
            .addConditionalEdges("differential_diagnosis_decision",
                    edge_async(new DiagnosisDispatcher()),
                    Map.of(
                            "needs_more_info", "additional_question", // 需要更多信息 -> 追加提问
                            "sufficient", "department_recommendation"  // 信息充足 -> 科室推荐
                    ))

            // 9. 追加提问 -> 患者补充信息 -> 疾病推理（重新推理）
            .addEdge("additional_question", "patient_info_supplementation")
            .addEdge("patient_info_supplementation", "disease_inference")

            // 10. 科室推荐 -> 风险复核
            .addEdge("department_recommendation", "risk_review")

            // 11. 风险复核后的分支
            .addConditionalEdges("risk_review",
                    edge_async(new RiskDispatcher()),
                    Map.of(
                            "high_risk", "emergency_triage",  // 高风险 -> 紧急分诊
                            "low_risk", "final_recommendation" // 低风险 -> 最终推荐
                    ))

            // 12. 紧急分诊 -> 最终推荐
            .addEdge("emergency_triage", "final_recommendation")

            // 结束节点
            .addEdge("final_recommendation", END)       // 问诊流程结束
            .addEdge("non_consult_response", END);      // 非问诊流程结束

    // 打印图形表示（PlantUML格式）
    GraphRepresentation graphRepresentation = stateGraph.getGraph(
            GraphRepresentation.Type.PLANTUML, "在线问诊工作流图示");
    logger.info("\n工作流图示:\n{}\n", graphRepresentation.content());

    return stateGraph;
  }

  // ========== 调度器实现 ==========

  /**
   * 意图分发器 - 根据意图判断结果决定流程走向
   */
  public static class IntentDispatcher implements EdgeAction {
    @Override
    public String apply(OverAllState state) {
      String intent = (String) state.value("intent").orElse("");
      logger.info("意图判断结果: {}", intent);

      // 更精确的意图判断
      if (intent.contains("medical_consultation")) {
        return "consultation";
      } else {
        return "non_consultation";
      }
    }
  }

  /**
   * 标准化分发器 - 判断是否需要人工审核
   */
  public static class StandardizationDispatcher implements EdgeAction {
    @Override
    public String apply(OverAllState state) {
      Boolean sensitiveFlag = (Boolean) state.value("sensitive_flag").orElse(false);
      logger.info("敏感词检查结果: {}", sensitiveFlag);

      return sensitiveFlag ? "sensitive" : "standardized";
    }
  }

  /**
   * 严重程度分发器 - 根据严重程度决定流程走向
   */
  public static class SeverityDispatcher implements EdgeAction {
    @Override
    public String apply(OverAllState state) {
      String severity = (String) state.value("severity").orElse("moderate");
      logger.info("症状严重程度: {}", severity);

      return severity.toLowerCase(); // 返回小写的严重程度分类
    }
  }

  /**
   * 诊断分发器 - 判断是否需要追加问题
   */
  public static class DiagnosisDispatcher implements EdgeAction {
    @Override
    public String apply(OverAllState state) {
      Boolean needAdditional = (Boolean) state.value("need_additional_questions").orElse(false);
      logger.info("是否需要追加问题: {}", needAdditional);

      return needAdditional ? "needs_more_info" : "sufficient";
    }
  }

  /**
   * 风险分发器 - 根据风险等级决定流程走向
   */
  public static class RiskDispatcher implements EdgeAction {
    @Override
    public String apply(OverAllState state) {
      String riskLevel = (String) state.value("risk_level").orElse("low_risk");
      logger.info("风险等级评估: {}", riskLevel);

      return riskLevel.toLowerCase(); // 返回小写的风险等级
    }
  }
}